#include <iostream>
#include <vector>
#include <string>

std::vector<std::string> g_vecJpgFiles;
void GetJpefFiles(void)
{
    const char* path = "/home/madengyun/MyGit/Cpp11Study/TestFfmpeg/build/";
    for (int i = 0; i < 300; i++)
    {
        char fileName[16] = {0};
        sprintf(fileName, "frame_%04d.jpg", i);
        std::string fullPath = std::string(path) + fileName;
        g_vecJpgFiles.push_back(fullPath);
    }
}

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <assert.h>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/opt.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
}


#define INPUT_FPS 25
#define OUTPUT_WIDTH 1920
#define OUTPUT_HEIGHT 1080

typedef struct {
    AVCodecContext *codec_ctx;
    AVFrame *frame;
    AVPacket *pkt;
    AVFormatContext *fmt_ctx;
    AVStream *stream;
    int frame_count;
} EncoderState;

// 初始化编码器和输出格式上下文
static int init_encoder(EncoderState *state, const char *filename) {
    // 查找H.264编码器
    const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        fprintf(stderr, "找不到H.264编码器\n");
        return -1;
    }

    // 分配编码器上下文
    state->codec_ctx = avcodec_alloc_context3(codec);
    if (!state->codec_ctx) {
        fprintf(stderr, "无法分配编码器上下文\n");
        return -1;
    }

    // 设置编码参数
    state->codec_ctx->width = OUTPUT_WIDTH;
    state->codec_ctx->height = OUTPUT_HEIGHT;
    state->codec_ctx->time_base = (AVRational){1, INPUT_FPS};
    state->codec_ctx->framerate = (AVRational){INPUT_FPS, 1};
    state->codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
    state->codec_ctx->bit_rate = 4000000;  // 4 Mbps
    state->codec_ctx->gop_size = 10;
    state->codec_ctx->max_b_frames = 1;

    // 打开编码器
    if (avcodec_open2(state->codec_ctx, codec, NULL) < 0) {
        fprintf(stderr, "无法打开编码器\n");
        return -1;
    }

    // 创建输出格式上下文
    avformat_alloc_output_context2(&state->fmt_ctx, NULL, NULL, filename);
    if (!state->fmt_ctx) {
        fprintf(stderr, "无法创建输出格式上下文\n");
        return -1;
    }

    // 添加视频流
    state->stream = avformat_new_stream(state->fmt_ctx, codec);
    if (!state->stream) {
        fprintf(stderr, "无法创建新的流\n");
        return -1;
    }

    state->stream->time_base = state->codec_ctx->time_base;
    
    // 复制编码器参数到流
    if (avcodec_parameters_from_context(state->stream->codecpar, state->codec_ctx) < 0) {
        fprintf(stderr, "无法复制编码器参数\n");
        return -1;
    }

    // 分配帧和包
    state->frame = av_frame_alloc();
    if (!state->frame) {
        fprintf(stderr, "无法分配帧\n");
        return -1;
    }

    state->frame->format = state->codec_ctx->pix_fmt;
    state->frame->width = state->codec_ctx->width;
    state->frame->height = state->codec_ctx->height;

    if (av_frame_get_buffer(state->frame, 0) < 0) {
        fprintf(stderr, "无法分配帧缓冲\n");
        return -1;
    }

    state->pkt = av_packet_alloc();
    if (!state->pkt) {
        fprintf(stderr, "无法分配包\n");
        return -1;
    }

    // 打开输出文件
    if (!(state->fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&state->fmt_ctx->pb, filename, AVIO_FLAG_WRITE) < 0) {
            fprintf(stderr, "无法打开输出文件\n");
            return -1;
        }
    }

    // 写入文件头
    if (avformat_write_header(state->fmt_ctx, NULL) < 0) {
        fprintf(stderr, "无法写入文件头\n");
        return -1;
    }

    return 0;
}

// 编码单帧
static int encode_frame(EncoderState *state) {
    int ret = avcodec_send_frame(state->codec_ctx, state->frame);
    if (ret < 0) {
        fprintf(stderr, "发送帧失败\n");
        return ret;
    }

    while (ret >= 0) {
        ret = avcodec_receive_packet(state->codec_ctx, state->pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            return 0;
        } else if (ret < 0) {
            fprintf(stderr, "编码错误\n");
            return ret;
        }

        // 设置正确的时间戳
        state->pkt->stream_index = state->stream->index;
        av_packet_rescale_ts(state->pkt, state->codec_ctx->time_base, state->stream->time_base);

        // 写入包到输出文件
        ret = av_interleaved_write_frame(state->fmt_ctx, state->pkt);
        if (ret < 0) {
            fprintf(stderr, "写入帧失败\n");
            return ret;
        }

        av_packet_unref(state->pkt);
    }

    return 0;
}

// 读取JPEG并转换为YUV帧
static int read_jpeg_to_frame(const char *filename, EncoderState *state) {
    AVFormatContext *input_ctx = NULL;
    if (avformat_open_input(&input_ctx, filename, NULL, NULL) < 0) {
        fprintf(stderr, "无法打开输入文件: %s\n", filename);
        return -1;
    }

    if (avformat_find_stream_info(input_ctx, NULL) < 0) {
        fprintf(stderr, "无法找到流信息\n");
        avformat_close_input(&input_ctx);
        return -1;
    }

    const AVCodec *jpeg_decoder = avcodec_find_decoder(AV_CODEC_ID_MJPEG);
    if (!jpeg_decoder) {
        fprintf(stderr, "找不到JPEG解码器\n");
        avformat_close_input(&input_ctx);
        return -1;
    }

    AVCodecContext *jpeg_ctx = avcodec_alloc_context3(jpeg_decoder);
    if (!jpeg_ctx) {
        fprintf(stderr, "无法分配JPEG解码器上下文\n");
        avformat_close_input(&input_ctx);
        return -1;
    }

    if (avcodec_parameters_to_context(jpeg_ctx, input_ctx->streams[0]->codecpar) < 0) {
        fprintf(stderr, "无法设置JPEG解码器参数\n");
        avcodec_free_context(&jpeg_ctx);
        avformat_close_input(&input_ctx);
        return -1;
    }

    if (avcodec_open2(jpeg_ctx, jpeg_decoder, NULL) < 0) {
        fprintf(stderr, "无法打开JPEG解码器\n");
        avcodec_free_context(&jpeg_ctx);
        avformat_close_input(&input_ctx);
        return -1;
    }

    AVFrame *jpeg_frame = av_frame_alloc();
    AVPacket *jpeg_packet = av_packet_alloc();

    struct SwsContext *sws_ctx = sws_getContext(
        jpeg_ctx->width, jpeg_ctx->height, jpeg_ctx->pix_fmt,
        OUTPUT_WIDTH, OUTPUT_HEIGHT, AV_PIX_FMT_YUV420P,
        SWS_BILINEAR, NULL, NULL, NULL
    );

    while (av_read_frame(input_ctx, jpeg_packet) >= 0) {
        if (avcodec_send_packet(jpeg_ctx, jpeg_packet) < 0) {
            fprintf(stderr, "发送JPEG包失败\n");
            break;
        }

        int ret = avcodec_receive_frame(jpeg_ctx, jpeg_frame);
        if (ret >= 0) {
            sws_scale(sws_ctx, 
                (const uint8_t * const *)jpeg_frame->data, jpeg_frame->linesize,
                0, jpeg_ctx->height,
                state->frame->data, state->frame->linesize
            );

            state->frame->pts = state->frame_count++;
            encode_frame(state);
        }

        av_packet_unref(jpeg_packet);
    }

    sws_freeContext(sws_ctx);
    av_frame_free(&jpeg_frame);
    av_packet_free(&jpeg_packet);
    avcodec_free_context(&jpeg_ctx);
    avformat_close_input(&input_ctx);

    return 0;
}

// 清理编码器资源
static void cleanup_encoder(EncoderState *state) {
    // Flush encoder
    avcodec_send_frame(state->codec_ctx, NULL);
    encode_frame(state);

    // 写入文件尾
    av_write_trailer(state->fmt_ctx);

    // 关闭输出文件
    if (state->fmt_ctx && !(state->fmt_ctx->oformat->flags & AVFMT_NOFILE))
        avio_closep(&state->fmt_ctx->pb);

    avcodec_free_context(&state->codec_ctx);
    av_frame_free(&state->frame);
    av_packet_free(&state->pkt);
    avformat_free_context(state->fmt_ctx);
}

int main(int argc, char *argv[]) {
    EncoderState state = {0};
    const char* output_filename = "output.mp4";

    if (init_encoder(&state, output_filename) < 0) {
        fprintf(stderr, "初始化编码器失败\n");
        cleanup_encoder(&state);
        return 1;
    }

    GetJpefFiles();

    // 处理所有输入的JPEG文件
    int jpeg_count = g_vecJpgFiles.size();
    for (int i = 0; i < jpeg_count; i++) {
        if (read_jpeg_to_frame(g_vecJpgFiles[i].c_str(), &state) < 0) {
            fprintf(stderr, "处理文件失败: %s\n", g_vecJpgFiles[i].c_str());
            continue;
        }
    }

    cleanup_encoder(&state);
    return 0;
}

